# -*- coding: utf-8 -*-
# @Time    : 2020/11/18 9:10
# @Author  : DuJiabao
# @File    : covert_tf2tflite.py
# @Description : 
import tensorflow as tf


def convert_tf2tflite(model_path: str):
    """
    将tf.saved_model.save保存的模型，转化成tflite格式的模型，以供树莓派使用。
    输出文件名：model_path+".tflite"
    :param model_path: 待转化的模型，一个model文件夹
    :return: 无
    """
    converter = tf.lite.TFLiteConverter.from_saved_model(model_path)
    tflite_model = converter.convert()

    output_path = model_path + ".tflite"

    # Save the model.
    with open(output_path, 'wb') as f:
        f.write(tflite_model)
    print(f'Have save tflite file in "{output_path}"')


if __name__ == '__main__':
    convert_tf2tflite("mnist_model")
